package main

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strconv"
	"strings"
)

type Proxy struct {
	reg    *regexp.Regexp
	remote *url.URL
	proxy  *httputil.ReverseProxy
}

func NewProxy(from string, remote *url.URL) *Proxy {
	var proxy = &Proxy{
		remote: remote,
	}
	if from != "" {
		proxy.reg = regexp.MustCompile(from)
	}
	proxy.proxy = &httputil.ReverseProxy{
		Director: proxy.requestFactory(),
		//ModifyResponse: responseFactory(value, remote, reg),
	}
	return proxy
}

func (t *Proxy) requestFactory() func(r *http.Request) {
	//remote := remotes[i]
	//reg := regs[i]
	return func(req *http.Request) {
		req.URL.Host = t.remote.Host
		if t.reg == nil {
			req.URL.Path = t.remote.Path
		} else {
			req.URL.Path = t.remote.Path + t.reg.ReplaceAllString(req.URL.Path, "")
		}
		req.URL.Scheme = t.remote.Scheme
		req.Host = t.remote.Host
		//req.Header.Set("Server", remote.Host)
		//req.Header.Set("Host", t.remote.Host)
		req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))

		//var contentType = getContentType(req)
		//if req.Method == "POST" && strings.Index(contentType, "multipart/formdata") == -1 {
		//	_,err := ioutil.ReadAll(req.Body)
		//	if err != nil {
		//		return
		//	}
		//	//可以加密
		//	//req.URL.Host = "127.0.0.1:8083"
		//	//req.Host = req.URL.Host
		//	//req.URL.Path = "/404"
		//	//req.Method = "GET"
		//	req.Header.Set("Content-Length", "0")
		//	req.Header.Set("content-length", "0")
		//	req.Body = ioutil.NopCloser(bytes.NewReader([]byte{}))
		//}
	}
}

func (t *Proxy) canProxy(url string) bool {
	if t.reg == nil {
		return false
	}
	return t.reg.MatchString(url)
}

func (t *Proxy) doProxy(w http.ResponseWriter, r *http.Request) {
	var cType = getContentType(r)
	var cLen = getContentLength(r)
	//如果不在加密列表中，直接代理
	if _, ok := ENCRYPT_CONTENT_TYPE_MAP[cType]; !ok {
		t.serveWithoutEncryptBody(w, r)
		return
	}
	//如果大小超过了限制
	if ENCRYPT_MAX_SIZE > 0 && cLen > ENCRYPT_MAX_SIZE {
		t.serveWithoutEncryptBody(w, r)
		return
	}
	//如果可以加密
	if r.Method != "POST" {
		t.serveWithoutEncryptBody(w, r)
		return
	}
	//加密body
	defer r.Body.Close()
	var bs = make([]byte, cLen)
	var read = 0
	var bufSize = 2048
	for read < cLen {
		n, err := r.Body.Read(bs[read:min(cLen, read+bufSize)])
		if err == io.EOF {
			break
		}
		if err != nil {
			http.NotFound(w, r)
			return
		}
		read += n
	}
	nbs, err := RsaEncrypt(bs, PUB_KEY)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(nbs))
	for k, v := range r.Header {
		if strings.EqualFold(k, "content-length") {
			r.Header.Set(k, strconv.Itoa(len(nbs)))
			r.ContentLength = int64(len(nbs))
		}
		if strings.EqualFold(k, "content-type") {
			r.Header.Set(k, "application/encryption; "+strings.Join(v, "; "))
		}
	}
	t.serveWithoutEncryptBody(w, r)
}

func (t *Proxy) serveWithoutEncryptBody(w http.ResponseWriter, r *http.Request) {
	t.proxy.ServeHTTP(w, r)
}

func getContentLength(r *http.Request) int {
	var length = 0
	var err error
	for k, v := range r.Header {
		if strings.EqualFold(k, "content-length") {
			length, err = strconv.Atoi(v[0])
			if err != nil {
				break
			}
		}
	}
	return length
}

func getContentType(r *http.Request) string {
	for k, v := range r.Header {
		if strings.EqualFold(k, "content-type") {
			return v[0]
		}
	}
	return ""
}

func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}
